from dataclasses import dataclass
from enum import Enum
from typing import Dict, List, Optional, Tuple, Union

from pathlib import Path
import numpy as np

import torch
from torch.utils.data import Dataset as TorchDataset

from sslsv.datasets.Sampler import SamplerConfig
from sslsv.datasets.DataAugmentation import DataAugmentation, DataAugmentationConfig
from sslsv.datasets.utils import load_audio


class FrameSamplingEnum(Enum):
    """
    Enumeration for frame sampling options for self-supervised training.

    Members:
        DEFAULT (str): Default frame sampling option.
        DINO (str): DINO frame sampling option.
    """

    DEFAULT = "default"
    DINO = "dino"


@dataclass
class DatasetConfig:
    """
    Dataset configuration.

    Attributes:
        ssl (bool): Whether to use SSLDataset for self-supervised learning (siamese).
        ssl_dino_global_count (int): Number of DINO global (large) frames.
        ssl_dino_global_length (int): Length of DINO global (large) frames.
        ssl_dino_local_count (int): Number of DINO local (small) frames.
        ssl_dino_local_length (int): Length of DINO local (small) frames.
        ssps (bool): Whether to use Self-Supervised Positive Sampling (SSPS).
        ssps_frame_length (int): Frame length for SSPS.
        augmentation (DataAugmentationConfig): Data-augmentation configuration.
        sampler (SamplerConfig): Sampler configuration.
        frame_length (int): Frame length (e.g. 32000 for 2s at 16kHz).
        frame_sampling (FrameSamplingEnum): Frame sampling option.
        max_samples (int): Maximum number of samples generated by the dataset.
        train (str): Path to train .csv file.
        label_key (str): Key of row to extract labels from train .csv file.
        base_path (Path): Base path for all files.
        num_workers (int): Number of workers for PyTorch DataLoader.
        pin_memory (bool): Whether to use pinned memory for PyTorch DataLoader.
    """

    ssl: bool = False
    ssl_dino_global_count: int = 2
    ssl_dino_global_length: int = 4 * 16000
    ssl_dino_local_count: int = 4
    ssl_dino_local_length: int = 2 * 16000
    ssps: bool = False
    ssps_frame_length: int = 64000
    augmentation: DataAugmentationConfig = None
    sampler: SamplerConfig = None
    frame_length: int = 32000
    frame_sampling: FrameSamplingEnum = FrameSamplingEnum.DEFAULT
    max_samples: int = None
    train: str = "voxceleb2_train.csv"
    label_key: str = "Speaker"
    base_path: Path = Path("./data/")
    num_workers: int = 8
    pin_memory: bool = True


class Dataset(TorchDataset):
    """
    PyTorch dataset generating audio frames.

    Attributes:
        config (DatasetConfig): Dataset configuration.
        files (List[str]): List of audio file paths.
        labels (Optional[List[int]]): List of labels.
        num_frames (int): Number of frames to sample.
        augmentation (DataAugmentation): Data-augmentation module (if enabled).
    """

    def __init__(
        self,
        config: DatasetConfig,
        files: List[str],
        labels: Optional[List[int]] = None,
        num_frames: int = 1,
    ):
        """
        Initialize a Dataset object.

        Args:
            config (DatasetConfig): Dataset configuration.
            files (List[str]): List of audio file paths.
            labels (Optional[List[int]]): List of labels. Defaults to None.
            num_frames (int): Number of frames to sample. Defaults to 1.

        Returns:
            None
        """
        super().__init__()

        self.config = config
        self.files = files
        self.labels = labels
        self.num_frames = num_frames

        self.augmentation = None
        if self.config.augmentation and self.config.augmentation.enable:
            self.augmentation = DataAugmentation(
                self.config.augmentation, self.config.base_path
            )

    def __len__(self) -> int:
        """
        Get the number of samples generated by the dataset.

        Returns:
            int: Length of the dataset.
        """
        if self.config.max_samples:
            return min(len(self.files), self.config.max_samples)
        return len(self.files)

    def preprocess_data(self, data: np.ndarray, augment: bool = True) -> np.ndarray:
        """
        Preprocess the input data by applying data-augmentation.

        Args:
            data (np.ndarray): Input audio. Shape: (1, L).
            augment (bool): Whether to apply data-augmentation. Defaults to True.

        Returns:
            np.ndarray: Output audio.

        Raises:
            AssertionError: If data is not a 2D array with shape (1, L).
        """
        if augment and self.augmentation:
            assert data.ndim == 2 and data.shape[0] == 1  # (1, L)
            data = self.augmentation(data)
        return data

    def __getitem__(
        self,
        i: int,
    ) -> Tuple[int, torch.Tensor, Dict[str, Union[str, int]]]:
        """
        Get an audio sample from the dataset.

        Args:
            i (int): Index of the sample.

        Returns:
            Tuple[int, torch.Tensor, Dict[str, Union[str, int]]]: Index, audio data,
                and additional info (file and label).
        """
        data = load_audio(
            self.config.base_path / self.files[i],
            frame_length=self.config.frame_length,
            num_frames=self.num_frames,
        )  # (N, T)

        x = torch.FloatTensor(self.preprocess_data(data)).squeeze(0)

        info = {"files": self.files[i]}
        if self.labels:
            info.update({"labels": self.labels[i]})

        return i, x, info
